# Copyright (c) Facebook, Inc. and its affiliates.
import logging
import os
import csv
import os.path as osp
from PIL import Image
import pandas as pd

from detectron2.utils.file_io import PathManager

logger = logging.getLogger(__name__)

def _get_dianjiao_files(path):
    file_list = []
    # scan through the directory
    if not isinstance(path, list):
        path = [path]
    for p in path:
        for root, dirs, files in os.walk(p):
            for f in files:
                if f.endswith('.jpg'):
                    file_list.append(os.path.join(root, f))
    assert len(file_list), "No images found in {}".format(path)
    assert PathManager.isfile(file_list[0]), f
    return file_list

def _get_csv_list(csv_path):
    pd_data = pd.read_csv(csv_path)
    image_list = [p.strip() for p in list(pd_data['img_pth'])]
    return image_list

def load_dianjiao_semantic(image_dir, gt_dir, save_csv=False):
    """
    Args:
        image_dir (str): path to the raw dataset. e.g., "~/dianjiao/leftImg8bit/train".
        gt_dir (str): path to the raw annotations. e.g., "~/dianjiao/gtFine/train".

    Returns:
        list[dict]: a list of dict, each has "file_name" and
            "sem_seg_file_name".
    """
    assert image_dir == gt_dir, "gt_file = image_file.replace('.jpg', '.png')"
    ret = []
    # gt_dir is small and contain many small files. make sense to fetch to local first
    if osp.basename(image_dir).endswith('.csv'):
        image_list = _get_csv_list(image_dir)
    else:
        gt_dir = PathManager.get_local_path(gt_dir)
        image_list = _get_dianjiao_files(image_dir)
    for image_file in image_list:
        label_file = image_file.replace(".jpg", ".png")
        if not os.path.exists(label_file):
            continue
        # mask = Image.open(label_file)
        # imgWidth, imgHeight = mask.size
        # if imgHeight > imgWidth:
        #     imgWidth, imgHeight = imgHeight, imgWidth
        ret.append(
            {
                "file_name": image_file,
                "sem_seg_file_name": label_file,
                # "height": imgHeight,
                # "width": imgWidth,
                "0_255_type": True
            }
        )
    assert len(ret), f"No images found in {image_dir}!"
    assert PathManager.isfile(
        ret[0]["sem_seg_file_name"]
    ), "Please generate labelTrainIds.png with cityscapesscripts/preparation/createTrainIdLabelImgs.py"  # noqa
    logger.info(f"find {len(image_list)} files, {len(ret)} is usable, {len(image_list) - len(ret)} have no annotations.")

    if save_csv and not osp.basename(image_dir).endswith('.csv'):
        save_path = osp.join(osp.dirname(gt_dir), osp.basename(gt_dir) + '.csv')
        with open(save_path, 'w', newline='') as csvfile:
            spamwriter = csv.writer(csvfile, delimiter=',',
                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
            spamwriter.writerow(['img_pth','img_name','source'] ) #

            for r in ret:
                p = r["file_name"]
                source = p.split('/')[len(gt_dir.split('/'))]
                img_name = osp.basename(p)
                spamwriter.writerow([p, img_name, source])
        csvfile.close()

    return ret
